{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[View the runnable example on GitHub](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/training/tensorflow/tensorflow_training_embedding_sparseadam.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Apply `SparseAdam` Optimizer for Large Embeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Embedding layers are often used to encode categorical items in deep learning applications. However, in applications such as recommendation systems, the embedding size may become huge due to large number of items or users, leading to extensive computational costs and space.\n",
    "\n",
    "For large embeddings, the batch size could be orders of magnitude smaller compared to the embedding matrix size. Thus, gradients to the embedding matrix in each batch could be sparse. Taking advantage of this, BigDL-Nano provides `bigdl.nano.tf.keras.layers.Embedding` and `bigdl.nano.tf.optimizers.SparseAdam` to accelerate large embeddings. `bigdl.nano.tf.optimizers.SparseAdam` is a variant of Adam which handles updates of sparse tensor more efficiently. `bigdl.nano.tf.keras.layers.Embedding` intends to avoid applying regularizer function directly to the embedding matrix, which further avoids making the sparse gradient dense."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "To apply Nano's `Embedding` layer and `SparseAdam` optimizer, you need to install BigDL-Nano for TensorFlow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "# install the nightly-built version of bigdl-nano for tensorflow;\n",
    "# intel-tensorflow will be installed at the meantime with intel's oneDNN optimizations enabled by default\n",
    "!pip install --pre --upgrade bigdl-nano[tensorflow]\n",
    "!source bigdl-nano-init  # set environment variables"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> Before starting your TensorFlow Keras application, it is highly recommended to run `source bigdl-nano-init` to set several environment variables based on your current hardware. Empirically, these variables will bring big performance increase for most TensorFlow Keras applications on training workloads."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "nbsphinx": "hidden"
   },
   "source": [
    "> ⚠️ **Warning**\n",
    "> \n",
    "> For Jupyter Notebook users, we recommend to run the commands above, especially `source bigdl-nano-init` before jupyter kernel is started, or some of the optimizations may not take effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "# install dependency for the dataset used in the following example\n",
    "!pip install tensorflow-datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To optimize your model for large embedding, you need to **import Nano's** `Embedding` **and** `SparseAdam` **first:**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bigdl.nano.tf.keras.layers import Embedding\n",
    "from bigdl.nano.tf.optimizers import SparseAdam\n",
    "\n",
    "# from tf.keras import Model\n",
    "from bigdl.nano.tf.keras import Model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> You could import `Model`/`Sequential` from `bigdl.nano.tf.keras` instead of `tf.keras` to gain more optimizations from Nano. Please refer to [API documentation](https://bigdl.readthedocs.io/en/latest/doc/PythonAPI/Nano/tensorflow.html#bigdl-nano-tf-keras) for more information."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take the [imdb_reviews](https://www.tensorflow.org/datasets/catalog/imdb_reviews) dataset as an example, and suppose we would like to train a model to classify movie reviews as positive/negative. Assuming that the vocabulary size of reviews is $20000$, and we want to fix the word vector to a length of $128$, we would have a big embedding matrix with size $20000 \\times 128$.\n",
    "\n",
    "To prepare the data for training, we need to process the samples as sequences of positive integers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "import re\n",
    "import string\n",
    "import tensorflow as tf\n",
    "import tensorflow_datasets as tfds\n",
    "from tensorflow.keras.layers import TextVectorization\n",
    "\n",
    "def create_datasets():\n",
    "    (raw_train_ds, raw_val_ds, raw_test_ds), info = tfds.load(\n",
    "        \"imdb_reviews\",\n",
    "        data_dir=\"/tmp/data\",\n",
    "        split=['train[:80%]', 'train[80%:]', 'test'],\n",
    "        as_supervised=True,\n",
    "        batch_size=32,\n",
    "        with_info=True\n",
    "    )\n",
    "\n",
    "    def custom_standardization(input_data):\n",
    "        lowercase = tf.strings.lower(input_data)\n",
    "        stripped_html = tf.strings.regex_replace(lowercase, \"<br />\", \" \")\n",
    "        return tf.strings.regex_replace(\n",
    "            stripped_html, f\"[{re.escape(string.punctuation)}]\", \"\"\n",
    "        )\n",
    "\n",
    "    vectorize_layer = TextVectorization(\n",
    "        standardize=custom_standardization,\n",
    "        max_tokens=20000,\n",
    "        output_mode=\"int\",\n",
    "        output_sequence_length=500,\n",
    "    )\n",
    "    \n",
    "    text_ds = raw_train_ds.map(lambda x, y: x)\n",
    "    vectorize_layer.adapt(text_ds)\n",
    "\n",
    "    def vectorize_text(text, label):\n",
    "        text = tf.expand_dims(text, -1)\n",
    "        return vectorize_layer(text), label\n",
    "\n",
    "    # vectorize the data\n",
    "    train_ds = raw_train_ds.map(vectorize_text)\n",
    "    val_ds = raw_val_ds.map(vectorize_text)\n",
    "    test_ds = raw_test_ds.map(vectorize_text)\n",
    "\n",
    "    return train_ds, val_ds, test_ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_ds, val_ds, test_ds = create_datasets()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; _The definition of_ `create_datasets` _can be found in the_ [runnable example](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/training/tensorflow/tensorflow_training_embedding_sparseadam.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We could then define the model. Same as using `tf.keras.layers.Embedding`, you could **instantiate a Nano's** `Embedding` **layer** as the first layer in the model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = tf.keras.Input(shape=(None,), dtype=\"int64\")\n",
    "\n",
    "# 20000 is the vocabulary size,\n",
    "# 128 is the embedding dimension\n",
    "x = Embedding(input_dim=20000, output_dim=128)(inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> If you would like to apply a regularizer function to the embedding matrix through setting `embeddings_regularizer`, Nano will apply the regularizer to the output tensors of the embedding layer instead to avoid making the sparse gradient dense (if `activity_regularize=None`).\n",
    ">\n",
    "> Please refer to [API document](https://bigdl.readthedocs.io/en/latest/doc/PythonAPI/Nano/tensorflow.html#bigdl.nano.tf.keras.layers.Embedding) for more information on `bigdl.nano.tf.keras.layers.Embedding`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, you could define the remaining parts of the model, and **configure the model for training with** `SparseAdam` **optimizer**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "nbsphinx": "hidden"
   },
   "outputs": [],
   "source": [
    "from tensorflow.keras import layers\n",
    "\n",
    "def make_backbone():\n",
    "    inputs = tf.keras.Input(shape=(None, 128))\n",
    "    x = layers.Dropout(0.5)(inputs)\n",
    "    x = layers.Conv1D(128, 7, padding=\"valid\", activation=\"relu\", strides=3)(x)\n",
    "    x = layers.Conv1D(128, 7, padding=\"valid\", activation=\"relu\", strides=3)(x)\n",
    "    x = layers.GlobalMaxPooling1D()(x)\n",
    "    x = layers.Dense(128, activation=\"relu\")(x)\n",
    "    x = layers.Dropout(0.5)(x)\n",
    "    predictions = layers.Dense(1, activation=\"sigmoid\", name=\"predictions\")(x)\n",
    "\n",
    "    model = Model(inputs, predictions)\n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# define the remaining layers of the model\n",
    "predictions = make_backbone()(x)\n",
    "model = Model(inputs, predictions)\n",
    "\n",
    "# Configure the model with Nano's SparseAdam optimizer\n",
    "model.compile(loss=\"binary_crossentropy\", optimizer=SparseAdam(), metrics=[\"accuracy\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "&nbsp;&nbsp;&nbsp;&nbsp;&nbsp; _The definition of_ `make_backbone` _can be found in the_ [runnable example](https://github.com/intel-analytics/BigDL/tree/main/python/nano/tutorial/notebook/training/tensorflow/tensorflow_training_embedding_sparseadam.ipynb)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📝 **Note**\n",
    ">\n",
    "> `SparseAdam` optimizer is a variant of `tf.keras.optimizers.Adam`. This method only updates moments that show up in the gradient, and applies only those portions of gradient to the trainable variables.\n",
    ">\n",
    "> Please refer to [API document](https://bigdl.readthedocs.io/en/latest/doc/PythonAPI/Nano/tensorflow.html#bigdl.nano.tf.optimizers.SparseAdam) for more information on `bigdl.nano.tf.optimizers.SparseAdam`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You could then train and evaluate your model as normal:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.fit(train_ds, validation_data=val_ds, epochs=10)\n",
    "model.evaluate(test_ds)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> 📚 **Related Readings**\n",
    "> \n",
    "> - [How to install BigDL-Nano](https://bigdl.readthedocs.io/en/latest/doc/Nano/Overview/install.html)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.7.13 ('nano-tf': conda)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.7.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "402532f56d486e9f832908f31130bbdf12bd8cb099dfb226783aa2c6b1479100"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
